# Written by Bram Cohen
# see LICENSE.txt for license information

import time
from random import randrange, shuffle
import math
from BitTornado.clock import clock
try:
    True
except:
    True = 1
    False = 0

class Choker:
    def __init__(self, config, schedule, picker, done = lambda: False):
        self.config = config
        self.round_robin_period = config['round_robin_period']
        self.schedule = schedule

        self.picker = picker
        self.connections = []
        self.last_preferred = 0
        self.last_round_robin = clock()
        self.done = done
        self.super_seed = False
        self.paused = False
        schedule(self._round_robin, 5)

    def set_round_robin_period(self, x):
        self.round_robin_period = x

    def _round_robin(self):
        self.schedule(self._round_robin, 5)
        if self.super_seed:
            cons = range(len(self.connections))
            to_close = []
            count = self.config['min_uploads']-self.last_preferred
            if count > 0:   # optimization
                shuffle(cons)
            for c in cons:
                i = self.picker.next_have(self.connections[c], count > 0)
                if i is None:
                    continue
                if i < 0:
                    to_close.append(self.connections[c])
                    continue
                self.connections[c].send_have(i)
                count -= 1
            for c in to_close:
                c.close()
        if self.last_round_robin + self.round_robin_period < clock():
            self.last_round_robin = clock()
            for i in xrange(1, len(self.connections)):
                c = self.connections[i]
                u = c.get_upload()
                if u.is_choked() and u.is_interested():
                    self.connections = self.connections[i:] + self.connections[:i]
                    break
        self._rechoke()

####### GROUP_VOD_CODE #######        
    def have_almost_all_the_file(self, conn, percent):
        counter = 0
        num_of_pieces = self.picker.get_num_of_pieces()
        for i in range(0,num_of_pieces):
            if conn.get_download().have[i]:
                counter += 1
        if (((counter/float(num_of_pieces))*100) >= percent):
            return True
        return False
    
    def have_pieces_after_me(self, conn):
        current, end = self.picker.computeWindowsSize()
        end = self.picker.get_num_of_pieces()
        for i in range(current, end):
            dont_have_piece = (self.picker.has[i]==0)
            the_other_client_have_it = conn.get_download().have[i]
            if (dont_have_piece and the_other_client_have_it):
                return True        
        return False
####### GROUP_VOD_CODE #######

    def _rechoke(self):
        preferred = []
        maxuploads = self.config['max_uploads']
        if self.paused:
            for c in self.connections:
                c.get_upload().choke()
            return
        if maxuploads > 1:
            
            for c in self.connections:
                u = c.get_upload()

                if not u.is_interested():
                    continue
                ####### GROUP_VOD_CODE #######
                d = c.get_download()
                r = d.get_rate()
                #when playing_done put group_vod clients in preferred list
                if ((self.playing_done()) and (not c.get_is_vod_client())):
                    continue
                
                ################scenarios to check for not vod clients####################
                ###if ((self.have_almost_all_the_file(c, 95)) and (not c.get_is_vod_client())): #if c is not a vod client and it has almost all the file, it will do restart very soon so we don't need to be liked by it  
                ###    continue
                #if ((not self.have_pieces_after_me(c)) and (not c.get_is_vod_client())): # if c is not a vod client and it doesn't have pieces that i don't have after my playing point, i don't need it so it won't get unchoke 
                #    continue
                #   else:
                #      r = 10000 #make sure vod_client will be in the preferred list... 
                ####### GROUP_VOD_CODE #######
                #if  d.is_snubbed(): #don't snubb clients!
                #    continue
                preferred.append((-r, c))
            self.last_preferred = len(preferred)
            preferred.sort()
            del preferred[maxuploads-1:]
            preferred = [x[1] for x in preferred]
        count = len(preferred)
        hit = False
        to_unchoke = []
        ####### GROUP_VOD_CODE #######
        
        #replace non vod client

        #check if some group_vod client needs a chunk in immediate window and if so, unchoke it.
        preferred_index = count-1 #the last connection in preferred
        for conn in self.connections:
            if conn.get_is_vod_client():
                if self.is_interested_in_window(conn) and (conn not in preferred):
                    #print 'connection is is_interested_in_window \r'
                    
                    if len(preferred) == maxuploads: #filled all slots, replace the least preferred
                        if (preferred_index < 0): #no more slots to replace
                            break
                        else:
                            while preferred_index>=0 and preferred[preferred_index].get_is_vod_client():
                                preferred_index-=1  
                            if (preferred_index < 0): #no more slots to replace
                                break                
                            preferred[preferred_index]=conn #override least preferred BT client
                            preferred_index-=1              # update index
                    else: #len(preferred) < maxuploads
                        preferred.append(conn)  # we still have place for connections
                        count +=1               #update count
                    
     
        ####### GROUP_VOD_CODE #######
        for c in self.connections:
            u = c.get_upload()
            if c in preferred:
                to_unchoke.append(u)
            else:
                if count < maxuploads or not hit: #if we add less than maxuploads conn to preferred list or we didn't choose a random conn that is interested already
                    ####### GROUP_VOD_CODE #######
                    if u.is_interested():
                        to_unchoke.append(u)
                    ####### GROUP_VOD_CODE #######
                        count += 1
                        hit = True
                else:
                    u.choke();
        for uploder in to_unchoke:
            uploder.unchoke()
            
        
    ####### GROUP_VOD_CODE #######
    def i_have_later_pieces_that_conn_need(self, conn, index):
        num_of_pieces = self.picker.get_num_of_pieces()
        for i in range(index,num_of_pieces):
            the_other_client_doesnt_have_it = (not (conn.get_download().have[i]))
            i_have_piece = (not (self.picker.has[i]==0))
            if (i_have_piece and the_other_client_doesnt_have_it):
                return True
        return False
    
    def compute_gap_difference(self, connection):
        strmwatcher = self.streamwatcher
        gap_diff = 0
        if strmwatcher.gap <> 0:
            peer_order = connection.order
            self_order = int(self.config["order"])
            ##if self=4, peer=1 than gap_diff is positive i.e. peer has a fore on us!!! 
            gap_diff = (self_order - peer_order) * strmwatcher.gap * (((float(strmwatcher.prefetchT) / 100) * strmwatcher.toKbytes(strmwatcher.total))/strmwatcher.rate) 
        return gap_diff
    
    def original_is_interested_in_window(self, connection):
        uploader = connection.get_upload()
        return uploader.is_interested()
    
    def new_is_interested_in_window(self, connection):
        if (connection.order is None): #didn't get order yet
            return True
        uploader = connection.get_upload()
        
        gap_diff = self.compute_gap_difference(connection)
        orig = self.get_playing_point(gap_diff)
        if (orig<0):
            orig=0
        
        prefetch_size_in_chunks = int((float(self.streamwatcher.prefetchT)/100) * (self.streamwatcher.numOfPieces)) 
        next_prefetch_chunk = prefetch_size_in_chunks
        while next_prefetch_chunk <= orig:
            next_prefetch_chunk += prefetch_size_in_chunks
            
        if ((not uploader.is_interested()) or (not self.streamwatcher) or (not self.i_have_later_pieces_that_conn_need(connection, next_prefetch_chunk))):
            return False
        
            
        #print '############ order:', self.config['order'] ,'peer:' , connection.order,  ' ###### playing point:', self.playing_point,'->',orig, '\r'
        
        return True
    
    def is_interested_in_window(self, connection):
        curr_gap = int(self.config['gap']) #we can compute it by computing where the other vod client are but it easier this way
        if curr_gap==0:
            return self.new_is_interested_in_window(connection)
        else: #curr_gap==2
            return self.original_is_interested_in_window(connection)
        
    
    def get_playing_point(self, gap_diff = 0):
        strmwatcher = self.streamwatcher
        t = time.time() - strmwatcher.startTime
        t = t + gap_diff

        orig  =  math.ceil(((t - strmwatcher.delay) * strmwatcher.rate) / strmwatcher.toKbytes(strmwatcher.piece_size))
        
        #self.playing_point =  int(((t2 - strmwatcher.delay) * strmwatcher.rate) / strmwatcher.toKbytes(strmwatcher.piece_size))
        return orig
    
    def playing_done(self):
        if self.streamwatcher:
            orig = self.get_playing_point()
            prefetch_size_in_chunks = int((float(self.streamwatcher.prefetchT)/100) * (self.streamwatcher.numOfPieces))
            return (orig >= (self.streamwatcher.numOfPieces-prefetch_size_in_chunks)) or self.done()
        else:
            return self.done()

    ####### GROUP_VOD_CODE #######

    def connection_made(self, connection, p = None):
        if p is None:
            p = randrange(-2, len(self.connections) + 1)
        self.connections.insert(max(p, 0), connection)
        self._rechoke()

    def connection_lost(self, connection):
        self.connections.remove(connection)
        self.picker.lost_peer(connection)
        if connection.get_upload().is_interested() and not connection.get_upload().is_choked():
            self._rechoke()

    def interested(self, connection):
        if not connection.get_upload().is_choked():
            self._rechoke()

    def not_interested(self, connection):
        if not connection.get_upload().is_choked():
            self._rechoke()

    def set_super_seed(self):
        while self.connections:             # close all connections
            self.connections[0].close()
        self.picker.set_superseed()
        self.super_seed = True

    def pause(self, flag):
        self.paused = flag
        self._rechoke()
